// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/tool/trace/trace_replay_task_runner.h>

namespace kumo::pollux::tool::trace {

std::pair<std::shared_ptr<exec::Task>, RowVectorPtr> TraceReplayTaskRunner::run(
    bool copyResults) {
  auto cursor = exec::TaskCursor::create(cursorParams_);
  std::vector<RowVectorPtr> results;
  auto* task = cursor->task().get();
  addSplits(task);

  while (cursor->moveNext()) {
    results.push_back(cursor->current());
  }

  task->taskCompletionFuture().wait();

  if (copyResults) {
    return {cursor->task(), copy(results)};
  }

  return {cursor->task(), nullptr};
}

std::shared_ptr<RowVector> TraceReplayTaskRunner::copy(
    const std::vector<RowVectorPtr>& results) {
  auto totalRows = 0;
  for (const auto& result : results) {
    totalRows += result->size();
  }
  auto copyResult = BaseVector::create<RowVector>(
      cursorParams_.planNode->outputType(),
      totalRows,
      memory::traceMemoryPool());
  auto resultRowOffset = 0;
  for (const auto& result : results) {
    copyResult->copy(result.get(), resultRowOffset, 0, result->size());
    resultRowOffset += result->size();
  }
  return copyResult;
}

TraceReplayTaskRunner& TraceReplayTaskRunner::maxDrivers(int32_t maxDrivers) {
  cursorParams_.maxDrivers = maxDrivers;
  return *this;
}

TraceReplayTaskRunner& TraceReplayTaskRunner::spillDirectory(
    const std::string& dir) {
  cursorParams_.spillDirectory = dir;
  return *this;
}

TraceReplayTaskRunner& TraceReplayTaskRunner::splits(
    const core::PlanNodeId& planNodeId,
    std::vector<exec::Split> splits) {
  splits_[planNodeId] = std::move(splits);
  return *this;
}

void TraceReplayTaskRunner::addSplits(exec::Task* task) {
  for (auto& [nodeId, nodeSplits] : splits_) {
    for (auto& split : nodeSplits) {
      task->addSplit(nodeId, std::move(split));
    }
    task->noMoreSplits(nodeId);
  }
}

} // namespace kumo::pollux::tool::trace
